/*
 * TSPExample
 * Created on 07.05.2010 13:56:22
 * Copyright (C) 2010 Signal-COM
 * $Id: TSPExample.java,v 1.2 2010/05/28 08:57:48 tvm Exp $
 */
package com.asteros.crypto.utils.sign;

import ru.signalcom.crypto.cms.*;
import ru.signalcom.crypto.tsp.*;

import javax.security.auth.x500.X500Principal;
import java.io.*;
import java.math.BigInteger;
import java.net.HttpURLConnection;
import java.net.URL;
import java.security.MessageDigest;
import java.security.Signature;
import java.security.cert.*;
import java.util.*;

public class SignUtils {
	private static String hashAlgorithm = DigestAlgorithm.GOST3411_CP;

	/**
	 * ������ ������������ ������������ (SignedData) ���������.
	 * @param data ������������� ������.
	 * @return ����������� ���������.
	 * @throws Exception
	 */
	public byte[] sign(byte[] data, Collection<Signer> signers) throws Exception {
		System.out.println("Signing...");
		InputStream in = new ByteArrayInputStream(data);
		ByteArrayOutputStream out = new ByteArrayOutputStream();

		SignedDataGenerator generator = new SignedDataGenerator(out);

		generator.addSigners(signers);
		generator.setDetached(false);
		generator.setContentType(ContentType.DATA);

		OutputStream sigOut = generator.open();

		byte[] buf = new byte[1024];
		int len;
		while ((len = in.read(buf)) > 0) {
			sigOut.write(buf, 0, len);
		}
		generator.close();
		in.close();
		return out.toByteArray();
	}

	/**
	 * ������ ������������ ������ ������� ��� �������, ������������
	 * � ����������� (SignedData) ���������.
	 * @param signed ����������� ���������.
	 * @return ����� ����������� ���������.
	 * @throws Exception
	 */
	public byte[] timestamp(byte[] signed, String url) throws Exception {

		System.out.println("Signature time-stamping...");

		InputStream in = new ByteArrayInputStream(signed);
		SignedDataParser signedData = new SignedDataParser(in);
		signedData.process(false);
		Collection<SignerInfo> signerInfos = signedData.getSignerInfos();
		Iterator<SignerInfo> it = signerInfos.iterator();
		while (it.hasNext()) {
			SignerInfo signerInfo = it.next();

			TimeStampReq request = makeTimeStampRequest(
					hashAlgorithm, signerInfo.getSignature(), null);

			//TimeStampToken token = getTimeStampToken(url, request);
			TimeStampToken token = getTimeStamp(url, request).getTimeStampToken();
			signerInfo.addUnsignedAttribute(new Attribute(
					AttributeType.SIGNATURE_TIME_STAMP, token.getEncoded()));
		}

		in.reset();
		ByteArrayOutputStream out = new ByteArrayOutputStream();
		SignedDataReplacer replacer = new SignedDataReplacer(in, out);
		replacer.setSignerInfos(signerInfos);
		replacer.replace();
		in.close();
		out.close();
		return out.toByteArray();
	}

	/**
	 * ������ ������������ ������� �� ����� �������.
	 * @param hashAlgorithm ������������� ��������� �����������.
	 * @param data ������.
	 * @param reqPolicy ������������� �������� ������� ������� �������.
	 * @return ������ �� ����� �������.
	 * @throws Exception
	 */
	private TimeStampReq makeTimeStampRequest(
			String hashAlgorithm,
			byte[] data,
			String reqPolicy) throws Exception {

		byte[] digest = MessageDigest.getInstance(hashAlgorithm).digest(data);

		TimeStampReqGenerator reqGen = new TimeStampReqGenerator(
				new MessageImprint(hashAlgorithm, digest));
		reqGen.setCertReq(true);
		if (reqPolicy != null) {
			reqGen.setReqPolicy(reqPolicy);
		}
		return reqGen.generate();
	}

	/**
	 * ������ ��������� ������ ������� �� TSP-�������.
	 * @param url ����� �������.
	 * @param req ������ �� ����� �������.
	 * @return ����� �������.
	 * @throws Exception
	 */
	/*	private TimeStampToken getTimeStampToken(
			String url,
			TimeStampReq req) throws Exception {

		TimeStampResp respEncoded = getTimeStamp(url, req);
		return respEncoded.getTimeStampToken();
	}*/

	private TimeStampResp getTimeStamp(
			String url,
			TimeStampReq req) throws Exception {

		System.out.println("Getting time-stamp token...");
		System.out.println("Connecting to " + url + " ...");
		URL u = new URL(url);
		HttpURLConnection conn = (HttpURLConnection) u.openConnection();
		conn.setRequestMethod("POST");
		conn.setUseCaches(false);
		//		conn.setRequestProperty("Content-Type", "application/timestamp-query");//application/timestamp-query
		//		conn.setRequestProperty("Accept", "application/timestamp-reply");
		conn.setDoInput(true);
		conn.setDoOutput(true);

		System.out.println("Sending request...");
		BufferedOutputStream out = new BufferedOutputStream(conn.getOutputStream());
		out.write(req.getEncoded());
		out.close();

		System.out.println("Getting response...");
		if (conn.getResponseCode() != HttpURLConnection.HTTP_OK) {
			throw new IOException("http request rejected: " + conn.getResponseCode() + " " + conn.getResponseMessage());
		}
		if (!conn.getContentType().equals("application/timestamp-reply")) {
			throw new IOException("invalid content type: " + conn.getContentType());
		}

		TimeStampResp resp = new TimeStampResp(conn.getInputStream());
		switch (resp.getStatus()) {
		case PKIStatus.GRANTED:
		case PKIStatus.GRANTED_WITH_MODS:
			break;
		default:
			throw new TSPException("Time-stamp request rejected: " + resp.getFailInfo() +
					(resp.getStatusString() != null ? " [" + resp.getStatusString() + "]" : ""));
		}
		resp.getTimeStampToken().check(req);
		
/*		X509Certificate cert = resp.getTimeStampToken().getTSACertificate();
	    File cert_file = new File("c:\\projects\\1.cer");
	    byte[] encoded = cert.getEncoded();
	    FileOutputStream outStream = new FileOutputStream(cert_file);
	    
	    outStream.write(encoded);*/

        System.out.println("getAccuracy - " + resp.getTimeStampToken().getTSTInfo().getAccuracy());
        System.out.println("getPolicy - " + resp.getTimeStampToken().getTSTInfo().getPolicy());
		resp.getTimeStampToken().validate(null);
		
		System.out.println("Response validated");
		return resp;	
	}

	/**
	 * ������ �������� ������������ (SignedData) ���������.
	 * @param signedData ����������� ���������.
	 * @throws Exception
	 */
	public void verify(
			byte[] signedData,
			Set<TrustAnchor> trust,
			List<CertStore> certStores, 
			byte[] content) throws Exception {

		System.out.println("Signature(s) verifying...");
//		System.out.println(new StrinsignedData);
//		System.out.println(new String(signedData).toCharArray());
		
		InputStream in = new ByteArrayInputStream(signedData);
		SignedDataParser parser = new SignedDataParser(in);
		parser.process();
		certStores.add(parser.getCertificatesAndCRLs());
		in.close();

		Collection<SignerInfo> signerInfos = parser.getSignerInfos();
		Iterator<SignerInfo> it = signerInfos.iterator();
		while (it.hasNext()) {
			SignerInfo signerInfo = it.next();
			verifySignerInfo(signerInfo, trust, certStores, content);
		}
		parser.close();
	}

	/**
	 * ������ �������� ����� ������� (SignerInfo).
	 * ��������:
	 * <br>1) �������� ������� ��� ������;</br>
	 * <br>2) �������� �����������;</br>
	 * <br>3) �������� �������������� �������� (���� ����).</br>
	 * @param signerInfo ���� �������.
	 * @param trust ������ ���������� ������������.
	 * @param certStores ��������� ������������ � ������� ���������� ������������.
	 * @param content �������������� ������ ��� ��������
	 * @throws Exception
	 */
	private void verifySignerInfo(
			SignerInfo signerInfo,
			Set<TrustAnchor> trust,
			List<CertStore> certStores, 
			byte[] content) throws Exception {

		boolean timeStamped = false;

		X509Certificate cert = lookupCertificate(certStores,
				signerInfo.getIssuer(), signerInfo.getSerialNumber());

		System.out.println("Signature verifying...");
		System.out.println("..." + cert.getSubjectX500Principal().getName());
		
		if (!signerInfo.verify(cert)) {
			throw new CMSException("signature failure");
		}
		
/*		��������� �������� ������ � getSignature �� �������� =( 
 * 		��� �� getSignature != ������� ����������� ������� ����� Signature.inisign+update+sign
 * 
		System.out.println(new String(signerInfo.getSignature()));
		if (content != null) {
			checkSignature(cert, signerInfo.getSignature(), content);
		}
		*/
		
		Collection<Attribute> attrs = signerInfo.getUnsignedAttributes();
		if (!attrs.isEmpty()) {
			Iterator<Attribute> it2 = attrs.iterator();
			while (it2.hasNext()) {
				Attribute at = it2.next();
				if (at.getType().compareTo(AttributeType.SIGNATURE_TIME_STAMP) == 0) {
					System.out.println("Signature time-stamp was found...");
					timeStamped = true;
					TimeStampToken token = new TimeStampToken(at.getValue());
					checkMessageImprint(token.getTSTInfo().getMessageImprint(), signerInfo.getSignature());
					verifyTimeStampToken(token, trust, certStores);

					verifyCertificate(cert, trust, certStores, token.getTSTInfo().getDate(), true);
				}
			}
		}

		if (!timeStamped) {
			verifyCertificate(cert, trust, certStores, null, true);
		}
	}

	/**
	 * ����� �����������.
	 * @param certStores ��������� ������������.
	 * @param issuer ��� ��������.
	 * @param serial �������� �����.
	 * @return ����������.
	 * @throws CertStoreException
	 */
	private X509Certificate lookupCertificate(
			List<CertStore> certStores,
			X500Principal issuer,
			BigInteger serial) throws CertStoreException {

		X509CertSelector csel = new X509CertSelector();
		csel.setIssuer(issuer);
		csel.setSerialNumber(serial);

		Iterator<CertStore> it = certStores.iterator();
		while (it.hasNext()) {
			CertStore certStore = it.next();
			Collection col = certStore.getCertificates(csel);
			if (!col.isEmpty()) {
				return (X509Certificate) col.iterator().next();
			}
		}
		throw new CertStoreException("certificate not found");
	}

	private void checkSignature(byte[] signedData, byte[] content) throws Exception {
		// �������� ���� ��� ���������� �����������
		byte[] detached = detach(signedData);
		
	    if (!Arrays.equals(content, detached)) {
	        throw new RuntimeException("�������� ������ � ������ � ������� �����������");
	    }
	}
	
	private void checkSignature(X509Certificate cert, byte[] signature, byte[] content) throws Exception {
		// �������� ���� ��� ���������� �����������
		
		Signature sign = Signature.getInstance("GOST3411withECGOST3410", "SC");
		sign.initVerify(cert);
		sign.update(content);
		
//		byte[] digest = MessageDigest.getInstance("1.2.643.2.2.9").digest(content);
//		sign.update(digest);

		System.out.println(sign.verify(signature));
		
		if (!sign.verify(signature)) {
//			throw new RuntimeException("signature failure");
		}
	}

	/**
	 * ������ �������� ������������ ������ � ���-��������.
	 * @param mi ���-��������.
	 * @param data ������.
	 * @throws Exception
	 */
	private void checkMessageImprint(MessageImprint mi, byte[] data)
	throws Exception {
		String alg = DigestAlgorithm.getName(mi.getHashAlgorithm());
		MessageDigest md = MessageDigest.getInstance(alg);
		md.update(data);
		if (!Arrays.equals(md.digest(), mi.getHashedMessage())) {
			throw new TSPException("message imprint does not match data");
		}
	}

	/**
	 * ������ �������� ������ �������.
	 * @param token ����� �������.
	 * @param trust ������ ���������� ������������.
	 * @param certStores ��������� ������������.
	 * @throws Exception
	 */
	private void verifyTimeStampToken(
			TimeStampToken token,
			Set trust,
			List<CertStore> certStores)
	throws Exception {
		System.out.println("Time-stamp token verifying...");
//		token.validate(null);
		verifyCertificate(token.getTSACertificate(), trust, certStores, null, false);   // TODO: must be true
	}

	/**
	 * ������ �������� �����������.
	 * @param cert ����������.
	 * @param trust ������ ���������� ������������.
	 * @param certStores ��������� ������������ � ������� ���������� ������������.
	 * @return ��������� ��������.
	 * @throws Exception
	 */
	private PKIXCertPathBuilderResult verifyCertificate(
			X509Certificate cert,
			Set<TrustAnchor> trust,
			List<CertStore> certStores,
			Date date,
			boolean revocationEnabled) throws Exception {

		System.out.println("X.509 certificate verifying...");
		System.out.println("Subject: " + cert.getSubjectX500Principal());
		X509CertSelector csel = new X509CertSelector();
		csel.setCertificate(cert);

		PKIXBuilderParameters params = new PKIXBuilderParameters(trust, csel);
		params.setDate(date);
		params.setRevocationEnabled(revocationEnabled);
		params.setCertStores(certStores);

		CertPathBuilder builder = CertPathBuilder.getInstance("PKIX", "SC");
		PKIXCertPathBuilderResult result = (PKIXCertPathBuilderResult) builder.build(params);
		System.out.println("Trust anchor: " + result.getTrustAnchor().getTrustedCert().getIssuerX500Principal());
		return result;
	}

	/**
	 * ������ ������� ��������, ������������
	 * � ����������� (SignedData) ����������.
	 * @param signed ����������� ���������.
	 * @return ����� ����������� ���������.
	 * @throws Exception
	 */
	public byte[] merge(byte[] signed, byte[] signedAdd) throws Exception {

		System.out.println("Signature time-stamping...");
		
		byte[] detached = detach(signed);
		byte[] detachedAdd = detach(signedAdd);
	    if (!Arrays.equals(detached, detachedAdd)) {
	        throw new RuntimeException("�������� � ������������� ������ �����������");
	    }
	    // ������ signerInfos � CertificatesAndCRLs
		InputStream in = new ByteArrayInputStream(signed);
		SignedDataParser signedData = getSignedData(in);
		Collection<SignerInfo> signerInfos = signedData.getSignerInfos();
		CertStore CertificatesAndCRLs =signedData.getCertificatesAndCRLs();
		
		InputStream inAdd = new ByteArrayInputStream(signedAdd);
		SignedDataParser signedDataAdd = getSignedData(inAdd);
		Collection<SignerInfo> signerInfosAdd = signedDataAdd.getSignerInfos();
		CertStore CertificatesAndCRLsAdd = signedDataAdd.getCertificatesAndCRLs();
		
		// ������� signerInfos
		signerInfos.addAll(signerInfosAdd);
		
		// ������� CertificatesAndCRLs		
//		//CertificateFactory cf = CertificateFactory.getInstance("X.509", "CS");
		//CertStoresEnv certStoresEnv = new CertStoresEnv(cf);
		CollectionCertStoreParameters c1 = (CollectionCertStoreParameters) CertificatesAndCRLs.getCertStoreParameters();
		CollectionCertStoreParameters c2 = (CollectionCertStoreParameters) CertificatesAndCRLsAdd.getCertStoreParameters();
		Collection coll1 = c1.getCollection();
		Collection coll2 = c2.getCollection();
		coll1.addAll(coll2);

		CollectionCertStoreParameters c3 = new CollectionCertStoreParameters(coll1);
	//	CertStoreParameters params = new CollectionCertStoreParameters(c3);
		CertStore CertificatesAndCRLsRes = CertStore.getInstance("Collection", c3 , "SC" );
		
		in.reset();
		ByteArrayOutputStream out = new ByteArrayOutputStream();
		SignedDataReplacer replacer = new SignedDataReplacer(in, out);
		replacer.setSignerInfos(signerInfos);
		replacer.setCertificatesAndCRLs(CertificatesAndCRLsRes);
		replacer.replace();
		in.close();
		out.close();
		return out.toByteArray();
	}

	private SignedDataParser getSignedData(InputStream in) throws Exception {
		SignedDataParser signedData = new SignedDataParser(in);
		signedData.process(false);
		return signedData;
	}
	/*
	private Collection<SignerInfo> getSignerInfos(InputStream in) throws Exception {
		SignedDataParser signedData = new SignedDataParser(in);
		signedData.process(false);
		Collection<SignerInfo> signerInfos = signedData.getSignerInfos();
		return signerInfos;
	}*/
	
    /**
     * ������ ��������� ����������� ������ �� ��������.
     * @param signed ����������� ���������.
     * @return ����������� ������.
     * @throws Exception
     */
    public byte[] detach(byte[] signed) throws Exception {
        System.out.println("Data detaching...");
        SignedDataParser parser = new SignedDataParser(new ByteArrayInputStream(signed));
        InputStream content = parser.getContent(false);
        ByteArrayOutputStream bOut = new ByteArrayOutputStream();
        byte[] data = new byte[1024];
        int len;
        while ((len = content.read(data)) >= 0) {
            bOut.write(data, 0, len);
        }
        parser.close();
        return bOut.toByteArray();
    }
}
